/**
 * @file hs_cuda_elementwise.cuh
 * @author iam002
 * @brief 并行应用场景: 逐元素运算 elementwise（多入单出 MISO）
 * @version 1.0
 * @date 2024-10-15
 * 
 * @copyright Copyright (c) 2024
 * 
 */
#ifndef __HS_CUDA_ELEMENTWISE_H__
#define __HS_CUDA_ELEMENTWISE_H__

#include "hs_cuda_utils.cuh"
#include "hs_cuda_elemwise_mimo.cuh"

namespace hs
{

namespace cuda
{

namespace elemwise
{


/**************************************************************
调用链:
    Unary/Binary/Ternary (顶层封装)
        |
        |
    GenericLauncher
        |
        |___ launchKernel
                |
                |___ applyGeneric (__global__)
                        |
                        |___ applyPack
                                |
                                |___ Functor (由用户设计功能单元)


功能单元:
    1. 在 device 端重载()运算符, 函数参数为输入变量, 函数返回值为输出值, 另外注意使用const修饰;
       例如: __device__ float operator(int a, int b, int c) const;
    2. 功能单元中的每个类成员必须有明确确定大小, 不要定义功能单元的类成员为:
        (a) std::vector 等标准库模板类
        (b) 指向 host 端地址的指针或指针数组
       个人猜测: 当功能单元传入核函数时, 会执行 cudaMalloc 和 cudaMemcpy 的操作, 在 device 端上
       复制一个同样的功能单元, 由于是 cudaMemcpy 是按照指定大小进行值的复制, 因此不建议成员中定义
       指向主机端地址的指针哈!

功能单元实例:

template<typename T>
struct AddFunctor {
    __device__ T operator()(T x, T y) const
    {
        return x + y;
    }
};

**************************************************************/


/* 解包, 执行功能单元, 展开可变长输入参数 */
template<int pack_size, typename FunctorType, typename OutputType, typename... InputType>
__device__ void applyPack(const FunctorType& functor
    , Packed<pack_size, OutputType>& dst_pack
    , const Packed<pack_size, InputType>&... src_pack)
{
    #pragma unroll
    for (int i = 0; i < pack_size; ++i) {
        dst_pack.elem[i] = functor((src_pack.elem[i])...);
    }
}


/* 核函数, 请注意, 核函数形参涉及到数组和指针的传递时, 应该保证它们是在device端, 这真的很重要;
 * 另外, 核函数也不要使用引用.
 */
template<int pack_size, typename FunctorType, typename OutputType, typename... InputType>
__global__ void __launch_bounds__(KERNEL_BLOCK_SIZE)
applyGeneric(FunctorType functor
        , int num_pack
        , Packed<pack_size, OutputType>* dst_pack
        , const Packed<pack_size, InputType>*... src_pack
        , int num_tail
        , OutputType* p_out_tail
        , const InputType*... p_in_tail)
{
    const int global_tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (int i = global_tid; i < num_pack; i += blockDim.x * gridDim.x) {
        /* 处理包 */
        applyPack<pack_size, FunctorType, OutputType, InputType...>(
            functor, dst_pack[i], (src_pack[i])...
        );
    }
    if (global_tid < num_tail) {
        /* 处理尾部数据 */
        p_out_tail[global_tid] = functor((p_in_tail[global_tid])...);
    }
}


/* 准备工作, 用于启动核函数 */
template<size_t pack_size, typename FunctorType, typename OutputType, typename... InputType>
cudaError_t launchKernel(const FunctorType & functor, int N, OutputType* p_out, const InputType*... p_in)
{
    /* 计算打包数和尾部剩余数据个数 */
    const int num_pack = N / pack_size;
    const int tail_offset = num_pack * pack_size;
    const int num_tail = N - tail_offset;

    /* 计算线程块数 */
    int num_blocks;
    {
        // 不调用 hs_cuda_utils.cu
        cudaError_t err = GetNumBlocks(num_pack, &num_blocks);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 调用核函数 */
    applyGeneric<pack_size, FunctorType, OutputType, InputType...><<<num_blocks, KERNEL_BLOCK_SIZE, 0>>>(
        functor
        , num_pack
        , reinterpret_cast<Packed<pack_size, OutputType>*>(p_out)
        , (reinterpret_cast<const Packed<pack_size, InputType>*>(p_in))...
        , num_tail
        , (p_out + tail_offset)
        , (p_in + tail_offset)... 
    );

    return cudaPeekAtLastError();
}


/* 再次封装, 隐藏 pack_size */
template<typename FunctorType, typename OutputType, typename... InputType>
struct GenericLauncher
{
    static cudaError_t launch(const FunctorType & functor, int N, OutputType* p_out, const InputType*... p_in)
    {
        constexpr int pack_size = PackSize<OutputType, InputType...>();
        
        cudaError_t err;
        if (isAlignedForPack<pack_size, OutputType, InputType...>(p_out, p_in...)) {
            err = launchKernel<pack_size, FunctorType, OutputType, InputType...>(functor, N, p_out, p_in...);
        }
        else {
            err = launchKernel<1, FunctorType, OutputType, InputType...>(functor, N, p_out, p_in...);
        }
        return err;
    }
};



/**
 * @brief 一元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam OutputType 输出数据类型
 * @tparam InputType 输入数据类型
 * @param N 数组长度
 * @param p_in 输入数据指针
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputType>
inline cudaError_t Unary(int N, const InputType* p_in, const FunctorType & functor,  OutputType* p_out)
{  
    return GenericLauncher<FunctorType, OutputType, InputType>::launch(functor, N, p_out, p_in);
}


/**
 * @brief 二元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam OutputType  输出数据类型
 * @tparam InputTypeA  输入数据类型A
 * @tparam InputTypeB  输入数据类型B
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB>
inline cudaError_t Binary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB>::launch(functor, N, p_out, p_in_a, p_in_b);
}


/**
 * @brief 三元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam OutputType   输出数据类型
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB, typename InputTypeC>
inline cudaError_t Ternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB, InputTypeC>::launch(functor, N, p_out, p_in_a, p_in_b, p_in_c);
}


/**
 * @brief 四元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam OutputType   输出数据类型
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @tparam InputTypeD   输入数据类型D
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param p_in_d 输入数据指针D
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB, typename InputTypeC, typename InputTypeD>
inline cudaError_t Quaternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const InputTypeD* p_in_d, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB, InputTypeC, InputTypeD>::launch(functor, N, p_out, p_in_a, p_in_b, p_in_c, p_in_d);
}

} // end namespace elemwise

/**
 * @brief 将数组的索引作为功能单元的输入参数之一, 并参与计算
 * 
 */
namespace elemwiseIndex
{

/**************************************************************

调用链:
    Unary/Binary/Ternary (顶层封装)
        |
        |
    GenericLauncher
        |
        |___ launchKernel
                |
                |___ applyGeneric (__global__)
                        |
                        |___ applyPack
                                |
                                |___ Functor (由用户设计功能单元)


功能单元:
    1. 在 device 端重载()运算符, 函数参数为输入变量(第一个输入变量为索引值), 函数
       返回值为输出值, 另外注意使用const修饰;
       例如: 
           __device__ float operator(int ID, float a, float b, float c) const;
    2. 功能单元中的每个类成员必须有明确确定大小, 不要定义功能单元的类成员为:
        (a) std::vector 等标准库模板类
        (b) 指向 host 端地址的指针或指针数组
       个人猜测: 当功能单元传入核函数时, 会执行 cudaMalloc 和 cudaMemcpy 的操作, 
       在 device 端上复制一个同样的功能单元; 由于是 cudaMemcpy 是按照指定大小进行
       值的复制, 因此不建议成员定义为指向主机端地址的指针变量. 

参考功能算子:

(单进单出, 输入与输出大小一致)

template<typename DataType>
struct FunctorMeshgridX
{

public:
    FunctorMeshgridX(int size_x, int size_y, DataType origin_x, DataType scale_x)
        : m_size_x(size_x), m_size_y(size_y), m_origin_x(origin_x), m_scale_x(scale_x)
    { }

    // id 是输入数组的索引
    __device__ DataType operator()(int id, DataType tmp) const
    {
        int yid = id / m_size_x;
        int xid = id - yid * m_size_x;
        DataType res = m_origin_x + m_scale_x * xid;
        return res;
    }

private:
    int m_size_x;
    int m_size_y;
    DataType m_origin_x;
    DataType m_scale_x;
};

**************************************************************/

/* 解包, 执行功能单元, 展开可变长输入参数 */
template<int pack_size, typename FunctorType, typename OutputType, typename... InputType>
__device__ void applyPack(const FunctorType& functor
    , int pack_id
    , Packed<pack_size, OutputType>& dst_pack
    , const Packed<pack_size, InputType>&... src_pack)
{
    #pragma unroll
    for (int i = 0; i < pack_size; ++i) {
        dst_pack.elem[i] = functor(pack_id * pack_size + i, (src_pack.elem[i])...);
    }
}

/* 核函数, 请注意, 核函数形参涉及到数组和指针的传递时, 应该保证它们是在device端, 这真的很重要;
 * 另外, 核函数也不要使用引用.
 */
template<int pack_size, typename FunctorType, typename OutputType, typename... InputType>
__global__ void __launch_bounds__(KERNEL_BLOCK_SIZE)
applyGeneric(FunctorType functor
        , int num_pack
        , Packed<pack_size, OutputType>* dst_pack
        , const Packed<pack_size, InputType>*... src_pack
        , int num_tail
        , OutputType* p_out_tail
        , const InputType*... p_in_tail)
{
    const int global_tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (int i = global_tid; i < num_pack; i += blockDim.x * gridDim.x) {
        /* 处理包 */
        applyPack<pack_size, FunctorType, OutputType, InputType...>(
            functor, i, dst_pack[i], (src_pack[i])...
        );
    }
    if (global_tid < num_tail) {
        /* 处理尾部数据 */
        const int tail_offset = num_pack * pack_size;
        p_out_tail[global_tid] = functor(global_tid + tail_offset, (p_in_tail[global_tid])...);
    }
}

/* 准备工作, 用于启动核函数 */
template<size_t pack_size, typename FunctorType, typename OutputType, typename... InputType>
cudaError_t launchKernel(const FunctorType & functor, int N, OutputType* p_out, const InputType*... p_in)
{
    /* 计算打包数和尾部剩余数据个数 */
    const int num_pack = N / pack_size;
    const int tail_offset = num_pack * pack_size;
    const int num_tail = N - tail_offset;

    /* 计算线程块数 */
    int num_blocks;
    {
        // 不调用 hs_cuda_utils.cu
        cudaError_t err = GetNumBlocks(num_pack, &num_blocks);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 调用核函数 */
    applyGeneric<pack_size, FunctorType, OutputType, InputType...><<<num_blocks, KERNEL_BLOCK_SIZE, 0>>>(
        functor
        , num_pack
        , reinterpret_cast<Packed<pack_size, OutputType>*>(p_out)
        , (reinterpret_cast<const Packed<pack_size, InputType>*>(p_in))...
        , num_tail
        , (p_out + tail_offset)
        , (p_in + tail_offset)... 
    );

    return cudaPeekAtLastError();
}

/* 再次封装, 隐藏 pack_size */
template<typename FunctorType, typename OutputType, typename... InputType>
struct GenericLauncher
{
    static cudaError_t launch(const FunctorType & functor, int N, OutputType* p_out, const InputType*... p_in)
    {
        constexpr int pack_size = PackSize<OutputType, InputType...>();
        
        cudaError_t err;
        if (isAlignedForPack<pack_size, OutputType, InputType...>(p_out, p_in...)) {
            err = launchKernel<pack_size, FunctorType, OutputType, InputType...>(functor, N, p_out, p_in...);
        }
        else {
            err = launchKernel<1, FunctorType, OutputType, InputType...>(functor, N, p_out, p_in...);
        }
        return err;
    }
};


/**
 * @brief 一元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam OutputType 输出数据类型
 * @tparam InputType 输入数据类型
 * @param N 数组长度
 * @param p_in 输入数据指针
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputType>
inline cudaError_t Unary(int N, const InputType* p_in, const FunctorType & functor, OutputType* p_out)
{  
    return GenericLauncher<FunctorType, OutputType, InputType>::launch(functor, N, p_out, p_in);
}


/**
 * @brief 二元运算模板
 * 
 * @tparam FunctorType 重载()运算的结构体
 * @tparam OutputType  输出数据类型
 * @tparam InputTypeA  输入数据类型A
 * @tparam InputTypeB  输入数据类型B
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB>
inline cudaError_t Binary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB>::launch(functor, N, p_out, p_in_a, p_in_b);
}



/**
 * @brief 三元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam OutputType   输出数据类型
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB, typename InputTypeC>
inline cudaError_t Ternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB, InputTypeC>::launch(functor, N, p_out, p_in_a, p_in_b, p_in_c);
}


/**
 * @brief 四元运算模板
 * 
 * @tparam FunctorType  重载()运算的结构体
 * @tparam OutputType   输出数据类型
 * @tparam InputTypeA   输入数据类型A
 * @tparam InputTypeB   输入数据类型B
 * @tparam InputTypeC   输入数据类型C
 * @tparam InputTypeD   输入数据类型D
 * @param N 数组长度
 * @param p_in_a 输入数据指针A
 * @param p_in_b 输入数据指针B
 * @param p_in_c 输入数据指针C
 * @param p_in_d 输入数据指针D
 * @param functor 函数器
 * @param p_out 输出数据指针
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB, typename InputTypeC, typename InputTypeD>
inline cudaError_t Quaternary(int N, const InputTypeA* p_in_a, const InputTypeB* p_in_b, const InputTypeC* p_in_c, const InputTypeD* p_in_d, const FunctorType & functor, OutputType* p_out)
{
    return GenericLauncher<FunctorType, OutputType, InputTypeA, InputTypeB, InputTypeC, InputTypeD>::launch(functor, N, p_out, p_in_a, p_in_b, p_in_c, p_in_d);
}

} // end namespace elemwiseIndex

} // end namespace cuda

} // end namespace hs

#endif

